package main

import (
	"bufio"
	"bytes"
	"context"
	"fmt"
	"io"
	"log"
	"math/rand/v2"
	"net"
	"net/url"
	"strconv"
	"strings"
	"time"
)

const (
	MAX_BUFF_SIZE = 1024
)

var (
	proxyMap = map[string]func(ctx context.Context, rawUrl *url.URL, conn net.Conn){
		"/register": func(ctx context.Context, rawUrl *url.URL, conn net.Conn) {
			p, err := getProxyPassFromRawUrl(rawUrl)
			if err != nil {
				writeBodyToConn(conn, "error proxy pass")
				return
			}

			weight, _ := strconv.Atoi(rawUrl.Query().Get("weight"))
			clear := rawUrl.Query().Get("clear_other_weight")
			if p.host == "" || p.port == "" || weight == 0 {
				writeBodyToConn(conn, "invalid params")
				return
			}

			proxyPassMap[p] = weight

			if clear != "" {
				for k := range proxyPassMap {
					if p.equals(k) {
						continue
					}
					proxyPassMap[k] = 0
				}
			}

			writeBodyToConn(conn, "success")
			fmt.Println(proxyPassMap)
		},
		"/unregister": func(ctx context.Context, rawUrl *url.URL, conn net.Conn) {
			p, err := getProxyPassFromRawUrl(rawUrl)
			if err != nil {
				writeBodyToConn(conn, "error proxy pass")
				return
			}

			delete(proxyPassMap, p)
			fmt.Printf("%#v\n", proxyPassMap)
			writeBodyToConn(conn, "success")
		},
	}

	proxyPassMap = make(map[proxyPass]int, 10)
)

func writeBodyToConn(conn net.Conn, body string) {
	// 构建一个简单的响应
	response := "HTTP/1.1 200 OK\r\n"
	response += "Content-Type: text/plain\r\n"
	response += fmt.Sprintf("Content-Length: %d\r\n", len(body))
	response += "Connection: close\r\n"
	response += "\r\n"
	response += body

	// 写入响应
	_, err := conn.Write([]byte(response))
	if err != nil {
		fmt.Println("Error writing response:", err.Error())
		return
	}
}

func getProxyPassFromRawUrl(rawUrl *url.URL) (proxyPass, error) {
	proxyPassStr := rawUrl.Query().Get("proxy_pass")
	s := strings.SplitN(proxyPassStr, ":", 2)
	if len(s) != 2 {
		fmt.Println(proxyPassStr, s)
		return proxyPass{}, fmt.Errorf("invalid proxy_pass")
	}

	return proxyPass{s[0], s[1]}, nil
}

type proxyPass struct {
	host string
	port string
}

func (p proxyPass) equals(other proxyPass) bool {
	return p.host == other.host && p.port == other.port
}

// NOW only HTTP
func main() {
	l, err := net.Listen("tcp", "127.0.0.1:8080")
	if err != nil {
		panic(err)
	}
	for {
		conn, err := l.Accept()
		if err != nil {
			log.Println(err)
			continue
		}

		//log.Println("accept conn")
		go writeBodyToConn(conn, "ok")
		//go handleConn(conn)
	}
}

func getRawUrl(r *bufio.Reader) (string, error) {
	// 读取第一行请求数据
	s, err := r.ReadString('\n')
	if err != nil {
		log.Println("read string error: ", err)
		return "", err
	}

	location := strings.Split(s, " ")[1]
	return location, nil
}

func getProxyPass() *proxyPass {
	if len(proxyPassMap) == 0 {
		return nil
	}

	if len(proxyPassMap) == 1 {
		for p := range proxyPassMap {
			return &p
		}
	}

	sum := 0
	for _, weight := range proxyPassMap {
		sum += weight
	}

	n := rand.IntN(sum)

	var (
		p      proxyPass
		weight int
	)

	for p, weight = range proxyPassMap {
		if weight > n {
			return &p
		}

		n -= weight
	}

	return &p
}

// 处理链接
func handleConn(conn net.Conn) {
	log.Println("handleConn")

	defer conn.Close()
	var request = make([]byte, MAX_BUFF_SIZE)

	// 从conn中读取请求数据
	n, err := conn.Read(request)
	if err != nil {
		fmt.Println("read request error: ", err)
		return
	}

	reader := bytes.NewReader(request[:n])
	r := bufio.NewReader(reader)

	rawURL, err := getRawUrl(r)
	if err != nil {
		log.Println("get raw url error: ", err)
		return
	}

	uri, err := url.Parse(rawURL)
	if err != nil {
		log.Println("parse url error: ", err)
		return
	}

	if handler, ok := proxyMap[uri.Path]; ok {
		log.Println("go proxy handler")
		handler(context.Background(), uri, conn)
		return
	}

	log.Println("read proxy pass")

	// 建立到想要请求的服务端的连接
	p := getProxyPass()
	if p == nil {
		log.Println("proxy pass is nil")
		writeBodyToConn(conn, "error proxy pass")
		return
	}

	fmt.Println(p.host, p.port)
	c, err := net.DialTimeout("tcp", net.JoinHostPort(p.host, p.port), 10*time.Second)
	if err != nil {
		fmt.Println(err)
		writeBodyToConn(conn, err.Error())
		return
	}
	defer c.Close()

	// 将需要请求的数据转发一份
	_, err = c.Write(request)
	if err != nil {
		fmt.Println("write request error: ", err)
		return
	}

	var buff [512]byte
	for {
		n, err := c.Read(buff[:])
		if err != nil {
			if err == io.EOF {
				break
			}

			fmt.Println("read error", err)
			return
		}

		_, err = conn.Write(buff[:n])
		if err != nil {
			fmt.Println("write to client error: ", err)
			return
		}
	}
}
